package com.cy.ragbase.config;

import io.milvus.client.MilvusServiceClient;
import io.milvus.param.ConnectParam;
import io.milvus.param.IndexType;
import io.milvus.param.MetricType;
import io.milvus.param.collection.*;
import io.milvus.param.index.CreateIndexParam;
import io.milvus.common.clientenum.ConsistencyLevelEnum;
import io.milvus.grpc.DataType;
import jakarta.annotation.PostConstruct;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

@Slf4j
@Configuration
public class MilvusConfig {

    @Value("${milvus.host}")
    private String host;

    @Value("${milvus.port}")
    private int port;

    @Value("${milvus.collection-name}")
    private String collectionName;

    @Value("${milvus.codeblock-collection-name}")
    private String codeBlockCollectionName;

    @Value("${milvus.dimension}")
    private int dimension;

    @Bean
    public MilvusServiceClient milvusClient() {
        ConnectParam connectParam = ConnectParam.newBuilder()
                .withHost(host)
                .withPort(port)
                .build();

        MilvusServiceClient client = new MilvusServiceClient(connectParam);
        log.info("Connected to Milvus at {}:{}", host, port);
        return client;
    }

    @PostConstruct
    public void initializeCollections() {
        try {
            // 直接创建客户端实例，而不是调用Bean方法
            ConnectParam connectParam = ConnectParam.newBuilder()
                    .withHost(host)
                    .withPort(port)
                    .build();
            MilvusServiceClient client = new MilvusServiceClient(connectParam);

            // 初始化文档集合
            initializeDocumentCollection(client);
            
            // 初始化代码块集合
            initializeCodeBlockCollection(client);

            log.info("All Milvus collections are ready");
            
            // 关闭临时客户端
            client.close();
        } catch (Exception e) {
            log.error("Failed to initialize Milvus collections", e);
            throw new RuntimeException("Failed to initialize Milvus", e);
        }
    }

    /**
     * 初始化文档集合
     */
    private void initializeDocumentCollection(MilvusServiceClient client) {
        try {
            // 检查文档集合是否存在
            HasCollectionParam hasCollectionParam = HasCollectionParam.newBuilder()
                    .withCollectionName(collectionName)
                    .build();

            boolean hasCollection = client.hasCollection(hasCollectionParam).getData();

            if (!hasCollection) {
                createDocumentCollection(client);
                createDocumentIndex(client);
                loadDocumentCollection(client);
                log.info("Document collection '{}' created and initialized", collectionName);
            } else {
                log.info("Document collection '{}' already exists", collectionName);
            }
        } catch (Exception e) {
            log.error("Failed to initialize document collection: {}", collectionName, e);
            throw new RuntimeException("Failed to initialize document collection", e);
        }
    }

    /**
     * 初始化代码块集合
     */
    private void initializeCodeBlockCollection(MilvusServiceClient client) {
        try {
            // 检查代码块集合是否存在
            HasCollectionParam hasCollectionParam = HasCollectionParam.newBuilder()
                    .withCollectionName(codeBlockCollectionName)
                    .build();

            boolean hasCollection = client.hasCollection(hasCollectionParam).getData();

            if (!hasCollection) {
                createCodeBlockCollection(client);
                createCodeBlockIndex(client);
                loadCodeBlockCollection(client);
                log.info("Code block collection '{}' created and initialized", codeBlockCollectionName);
            } else {
                log.info("Code block collection '{}' already exists", codeBlockCollectionName);
            }
        } catch (Exception e) {
            log.error("Failed to initialize code block collection: {}", codeBlockCollectionName, e);
            throw new RuntimeException("Failed to initialize code block collection", e);
        }
    }

    /**
     * 创建文档集合
     */
    private void createDocumentCollection(MilvusServiceClient client) {
        // 定义字段
        FieldType idField = FieldType.newBuilder()
                .withName("id")
                .withDataType(DataType.VarChar)
                .withMaxLength(64)
                .withPrimaryKey(true)
                .withAutoID(false)
                .build();

        FieldType documentIdField = FieldType.newBuilder()
                .withName("document_id")
                .withDataType(DataType.VarChar)
                .withMaxLength(64)
                .build();

        FieldType chunkIndexField = FieldType.newBuilder()
                .withName("chunk_index")
                .withDataType(DataType.Int64)
                .build();

        FieldType contentField = FieldType.newBuilder()
                .withName("content")
                .withDataType(DataType.VarChar)
                .withMaxLength(2000)
                .build();

        FieldType vectorField = FieldType.newBuilder()
                .withName("vector")
                .withDataType(DataType.FloatVector)
                .withDimension(dimension)
                .build();

        // 创建集合
        CreateCollectionParam createCollectionParam = CreateCollectionParam.newBuilder()
                .withCollectionName(collectionName)
                .withDescription("RAG document vectors")
                .withShardsNum(2)
                .addFieldType(idField)
                .addFieldType(documentIdField)
                .addFieldType(chunkIndexField)
                .addFieldType(contentField)
                .addFieldType(vectorField)
                .withConsistencyLevel(ConsistencyLevelEnum.STRONG)
                .build();

        client.createCollection(createCollectionParam);
        log.info("Created document collection: {}", collectionName);
    }

    /**
     * 创建代码块集合
     */
    private void createCodeBlockCollection(MilvusServiceClient client) {
        // 定义代码块集合的字段
        FieldType idField = FieldType.newBuilder()
                .withName("id")
                .withDataType(DataType.VarChar)
                .withMaxLength(64)
                .withPrimaryKey(true)
                .withAutoID(false)
                .build();

        FieldType nameField = FieldType.newBuilder()
                .withName("name")
                .withDataType(DataType.VarChar)
                .withMaxLength(200)
                .build();

        FieldType codeField = FieldType.newBuilder()
                .withName("code")
                .withDataType(DataType.VarChar)
                .withMaxLength(20000)  // 代码内容可能比较长
                .build();

        FieldType infoField = FieldType.newBuilder()
                .withName("info")
                .withDataType(DataType.VarChar)
                .withMaxLength(1000)
                .build();

        FieldType languageField = FieldType.newBuilder()
                .withName("language")
                .withDataType(DataType.VarChar)
                .withMaxLength(50)
                .build();

        FieldType labelsField = FieldType.newBuilder()
                .withName("labels")
                .withDataType(DataType.VarChar)
                .withMaxLength(500)  // 存储逗号分隔的标签字符串
                .build();

        FieldType vectorField = FieldType.newBuilder()
                .withName("vector")
                .withDataType(DataType.FloatVector)
                .withDimension(dimension)
                .build();

        // 创建代码块集合
        CreateCollectionParam createCollectionParam = CreateCollectionParam.newBuilder()
                .withCollectionName(codeBlockCollectionName)
                .withDescription("RAG code block vectors")
                .withShardsNum(2)
                .addFieldType(idField)
                .addFieldType(nameField)
                .addFieldType(codeField)
                .addFieldType(infoField)
                .addFieldType(languageField)
                .addFieldType(labelsField)
                .addFieldType(vectorField)
                .withConsistencyLevel(ConsistencyLevelEnum.STRONG)
                .build();

        client.createCollection(createCollectionParam);
        log.info("Created code block collection: {}", codeBlockCollectionName);
    }

    /**
     * 为文档集合创建索引
     */
    private void createDocumentIndex(MilvusServiceClient client) {
        CreateIndexParam indexParam = CreateIndexParam.newBuilder()
                .withCollectionName(collectionName)
                .withFieldName("vector")
                .withIndexType(IndexType.IVF_FLAT)
                .withMetricType(MetricType.COSINE)
                .withExtraParam("{\"nlist\":128}")
                .build();

        client.createIndex(indexParam);
        log.info("Created index for document collection: {}", collectionName);
    }

    /**
     * 为代码块集合创建索引
     */
    private void createCodeBlockIndex(MilvusServiceClient client) {
        CreateIndexParam indexParam = CreateIndexParam.newBuilder()
                .withCollectionName(codeBlockCollectionName)
                .withFieldName("vector")
                .withIndexType(IndexType.IVF_FLAT)
                .withMetricType(MetricType.COSINE)
                .withExtraParam("{\"nlist\":128}")
                .build();

        client.createIndex(indexParam);
        log.info("Created index for code block collection: {}", codeBlockCollectionName);
    }

    /**
     * 加载文档集合
     */
    private void loadDocumentCollection(MilvusServiceClient client) {
        LoadCollectionParam loadParam = LoadCollectionParam.newBuilder()
                .withCollectionName(collectionName)
                .build();

        client.loadCollection(loadParam);
        log.info("Loaded document collection: {}", collectionName);
    }

    /**
     * 加载代码块集合
     */
    private void loadCodeBlockCollection(MilvusServiceClient client) {
        LoadCollectionParam loadParam = LoadCollectionParam.newBuilder()
                .withCollectionName(codeBlockCollectionName)
                .build();

        client.loadCollection(loadParam);
        log.info("Loaded code block collection: {}", codeBlockCollectionName);
    }
}